// Copyright (c) 2019, the R8 project authors. Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

package com.android.tools.r8.graph.analysis;

import com.android.tools.r8.cf.code.CfConstNumber;
import com.android.tools.r8.cf.code.CfFieldInstruction;
import com.android.tools.r8.cf.code.CfGoto;
import com.android.tools.r8.cf.code.CfIf;
import com.android.tools.r8.cf.code.CfInstruction;
import com.android.tools.r8.cf.code.CfInvoke;
import com.android.tools.r8.cf.code.CfLogicalBinop;
import com.android.tools.r8.graph.CfCode;
import com.android.tools.r8.graph.DexEncodedMethod;
import com.android.tools.r8.graph.DexItemFactory;
import com.android.tools.r8.ir.optimize.info.OptimizationFeedback;
import com.google.common.collect.ImmutableList;
import java.util.List;
import org.objectweb.asm.Opcodes;

public class ClassInitializerAssertionEnablingAnalysis extends EnqueuerAnalysis {
  private final DexItemFactory dexItemFactory;
  private final OptimizationFeedback feedback;

  public ClassInitializerAssertionEnablingAnalysis(
      DexItemFactory dexItemFactory, OptimizationFeedback feedback) {
    this.dexItemFactory = dexItemFactory;
    this.feedback = feedback;
  }

  @Override
  public void processNewlyLiveMethod(DexEncodedMethod method) {
    if (method.isClassInitializer()) {
      if (method.getCode().isCfCode()) {
        if (hasJavacClinitAssertionCode(method.getCode().asCfCode())) {
          feedback.setInitializerEnablingJavaAssertions(method);
        }
      }
    }
  }

  // The <clinit> instruction sequence generated by javac for classes which use assertions. The
  // call to desiredAssertionStatus() will provide the value for the -ea option for the class.
  //
  //  0: ldc           #<n>                // Current class
  //  2: invokevirtual #<n>                // Method java/lang/Class.desiredAssertionStatus:()Z
  //  5: ifne          12
  //  8: iconst_1
  //  9: goto          13
  // 12: iconst_0
  // 13: putstatic     #<n>                // Field $assertionsDisabled:Z

  // R8 processing with class file backend will rewrite this sequence to either of the following
  // depending on debug or release mode.

  //  2: ldc           #<n>                // Current class
  //  3: invokevirtual #<n>                // Method java/lang/Class.desiredAssertionStatus:()Z
  //  4: istore        0
  //  5: iload         0
  //  6: ldc           1
  //  7: ixor
  //  8: istore        0
  //  9: iload         0
  // 13: putstatic     #<n>                // Field $assertionsDisabled:Z

  //  2: ldc           #<n>                // Current class
  //  3: invokevirtual                     // Method java/lang/Class.desiredAssertionStatus:()Z
  //  4: ldc           1
  //  5: ixor
  // 13: putstatic     #<n>                // Field $assertionsDisabled:Z

  private static List<Class<?>> javacInstructionSequence =
      ImmutableList.of(
          CfIf.class,
          CfConstNumber.class,
          CfGoto.class,
          CfConstNumber.class,
          CfFieldInstruction.class);
  private static List<Class<?>> r8InstructionSequence =
      ImmutableList.of(CfConstNumber.class, CfLogicalBinop.class, CfFieldInstruction.class);

  private boolean hasJavacClinitAssertionCode(CfCode code) {
    for (int i = 0; i < code.instructions.size(); i++) {
      CfInstruction instruction = code.instructions.get(i);
      if (instruction.isInvoke()) {
        // Check for the generated instruction sequence by looking for the call to
        // desiredAssertionStatus() followed by the expected instruction types and finally checking
        // the exact target of the putstatic ending the sequence.
        CfInvoke invoke = instruction.asInvoke();
        if (invoke.getOpcode() == Opcodes.INVOKEVIRTUAL
            && invoke.getMethod() == dexItemFactory.classMethods.desiredAssertionStatus) {
          CfFieldInstruction fieldInstruction = isJavacInstructionSequence(code, i + 1);
          if (fieldInstruction == null) {
            fieldInstruction = isR8InstructionSequence(code, i + 1);
          }
          if (fieldInstruction != null) {
            return fieldInstruction.getOpcode() == Opcodes.PUTSTATIC
                && fieldInstruction.getField().name == dexItemFactory.assertionsDisabled;
          }
        }
      }
      // Only check initial straight line code.
      if (instruction.isJump()) {
        return false;
      }
    }
    return false;
  }

  private CfFieldInstruction isJavacInstructionSequence(CfCode code, int fromIndex) {
    List<Class<?>> sequence = javacInstructionSequence;
    int nextExpectedInstructionIndex = 0;
    CfInstruction instruction = null;
    for (int i = fromIndex;
        i < code.instructions.size() && nextExpectedInstructionIndex < sequence.size();
        i++) {
      instruction = code.instructions.get(i);
      if (instruction.isLabel()) {
        // Just ignore labels.
        continue;
      }
      if (instruction.getClass() != sequence.get(nextExpectedInstructionIndex)) {
        break;
      }
      nextExpectedInstructionIndex++;
    }
    return nextExpectedInstructionIndex == sequence.size()
        ? instruction.asFieldInstruction()
        : null;
  }

  private CfFieldInstruction isR8InstructionSequence(CfCode code, int fromIndex) {
    List<Class<?>> sequence = r8InstructionSequence;
    int nextExpectedInstructionIndex = 0;
    CfInstruction instruction = null;
    for (int i = fromIndex;
        i < code.instructions.size() && nextExpectedInstructionIndex < sequence.size();
        i++) {
      instruction = code.instructions.get(i);
      if (instruction.isStore() || instruction.isLoad()) {
        // Just ignore stores and loads.
        continue;
      }
      if (instruction.getClass() != sequence.get(nextExpectedInstructionIndex)) {
        break;
      }
      nextExpectedInstructionIndex++;
    }
    return nextExpectedInstructionIndex == sequence.size()
        ? instruction.asFieldInstruction()
        : null;
  }
}
